from sentence_transformers import SentenceTransformer
from pymilvus import connections, Collection, FieldSchema, CollectionSchema, DataType, utility

# 指定本地模型路径
model_path = r'E:\repository\models\paraphrase-multilingual-MiniLM-L12-v2'  # 替换为你的本地路径

# 加载本地模型
model = SentenceTransformer(model_path)
# 准备数据
documents = [
    {"id": 1, "text": "人工智能的发展与应用", "category": "technology"},
    {"id": 2, "text": "自然语言处理技术进展", "category": "technology"},
    {"id": 3, "text": "梦的解析", "category": "technology"},
    {"id": 4, "text": "时间简史", "category": "technology"},
    {"id": 5, "text": "历史简史", "category": "technology"},
    {"id": 5, "text": "高等数学", "category": "technology"},
    {"id": 5, "text": "线性代数", "category": "technology"},
    {"id": 5, "text": "微积分", "category": "technology"},
]
# 向量化
texts = [doc["text"] for doc in documents]
embeddings = model.encode(texts)
# 检查向量形状
print(f"Embeddings shape: {embeddings.shape}")  # 应该是 (5, 384)

# 插入数据
entities = [
    [doc["id"] for doc in documents],  # id列表
    embeddings.tolist(),               # 向量列表
    [doc["text"] for doc in documents], # 文本列表
    [doc["category"] for doc in documents] # 类别列表
]
collection_name = "documents"
connections.connect("default", host="localhost", port="19530")

# 检查集合是否存在
if utility.has_collection(collection_name):
    # 连接Milvus
    # 如果存在，直接加载集合
    collection = Collection(name=collection_name)
    print(f"Collection {collection_name} already exists. Using existing collection.")
    collection.load()  # 加载集合到内存
    collection.insert(entities)